package org.example.mqstream.offsettracking;

import com.rabbitmq.stream.ByteCapacity;
import com.rabbitmq.stream.Environment;
import com.rabbitmq.stream.OffsetSpecification;

import java.time.Duration;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import static java.nio.charset.StandardCharsets.UTF_8;

public class OffsetTrackingReceive {

    /**
     * 主程序入口
     * 本程序示例展示了如何使用Pulsar客户端API来消费消息，并进行偏移量的手动跟踪
     * 特别地，它演示了如何设置环境、创建流、定义消费策略以及处理消息
     *
     * @param args 命令行参数
     * @throws InterruptedException 如果在等待消息消费完成时线程被中断
     */
    public static void main(String[] args) throws InterruptedException {
        // 创建Pulsar环境配置，指定主机、端口、用户名、密码以及心跳间隔
        try (Environment environment =
                     Environment.builder()
                             .host("192.168.0.228")
                             .port(5552)
                             .username("mq")
                             .password("123456")
                             .requestedHeartbeat(Duration.ofSeconds(5))
                             .build()) {
            // 定义要操作的流名称，并创建一个最大长度为1GB的流
            String stream = "stream-offset-tracking-java";
            environment.streamCreator().stream(stream).maxLengthBytes(ByteCapacity.GB(1)).create();

            // 初始化偏移量规范为从流的起始位置开始消费
            OffsetSpecification offsetSpecification = OffsetSpecification.first();
            // 用于记录第一个偏移量、最后一个偏移量和消息计数的原子变量
            AtomicLong firstOffset = new AtomicLong(-1);
            AtomicLong lastOffset = new AtomicLong(0);
            AtomicLong messageCount = new AtomicLong(0);
            // 用于在消息消费完成后通知主线程的计数器
            CountDownLatch consumedLatch = new CountDownLatch(1);

            // 构建消费者，指定流、偏移量规范、消费名称和手动偏移量跟踪策略
            environment.consumerBuilder().stream(stream)
                    .offset(offsetSpecification)
                    .name("offset-tracking-tutorial")
                    .manualTrackingStrategy().builder()
                    .messageHandler(
                            (ctx, msg) -> {
                                // 当接收到第一个消息时，记录其偏移量
                                if (firstOffset.compareAndSet(-1, ctx.offset())) {
                                    System.out.println("First message received.");
                                }
                                // 每消费10条消息，存储一次偏移量
                                if (messageCount.incrementAndGet() % 10 == 0) {
                                    ctx.storeOffset();
                                }
                                // 解析消息体，如果遇到特定的结束标记，记录最后一个偏移量并关闭消费者
                                String body = new String(msg.getBodyAsBinary(), UTF_8);
                                if ("marker".equals(body)) {
                                    lastOffset.set(ctx.offset());
                                    ctx.storeOffset();
                                    ctx.consumer().close();
                                    consumedLatch.countDown();
                                }
                            })
                    .build();
            // 启动消费者
            System.out.println("Started consuming...");

            // 等待消息消费完成或60分钟超时
            consumedLatch.await(60, TimeUnit.MINUTES);

            // 打印消费完成时的首个和最后一个偏移量
            System.out.printf("Done consuming, first offset %d, last offset %d.%n",
                    firstOffset.get(), lastOffset.get());
        }
    }


}